package com.chatplus.application.aiprocessor.handler;

import com.chatplus.application.aiprocessor.platform.chat.ChatProcessorService;
import com.chatplus.application.aiprocessor.provider.ChatAiProcessorServiceProvider;
import com.chatplus.application.domain.dto.ws.ChatRecordMessage;
import com.chatplus.application.enumeration.AiPlatformEnum;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

@Component
public class PromptHandler {
    private final ChatAiProcessorServiceProvider chatAiProcessorServiceProvider;


    public PromptHandler(ChatAiProcessorServiceProvider chatAiProcessorServiceProvider) {
        this.chatAiProcessorServiceProvider = chatAiProcessorServiceProvider;
    }

    public String rewrite(String prompt) {
        String rewritePromptTemplate = "Please rewrite the following text into AI painting prompt words, and please try to add detailed description of the picture, painting style, scene, rendering effect, picture light and other elements. Please output directly in English without any explanation, within 150 words. The text to be rewritten is: [%s]";
        return action(rewritePromptTemplate, prompt);
    }

    public String translateToEN(String prompt) {
        String translatePromptTemplateEn = "Translate the following painting prompt words into English keyword phrases. Without any explanation, directly output the keyword phrases separated by commas. The content to be translated is: [%s]";
        return action(translatePromptTemplateEn, prompt);
    }

    public String translateToCN(String prompt) {
        String translatePromptTemplateCn = "Translate the following prompt words into Chinese. Without any explanation, directly output. The content to be translated is: [%s]";
        return action(translatePromptTemplateCn, prompt);
    }

    private String action(String promptTemplate, String prompt) {
        ChatProcessorService chatProcessorService = chatAiProcessorServiceProvider.getAiProcessorServiceByPlatform(AiPlatformEnum.OPEN_AI);
        List<ChatRecordMessage> messageContext = new ArrayList<>();
        ChatRecordMessage chatRecordMessage = ChatRecordMessage.builder().prompt(String.format(promptTemplate, prompt)).build();
        messageContext.add(chatRecordMessage);
        // TODO 加一个特殊模型的类型
//        ChatModelEntity chatModelEntity = ChatModelEntity.builder().value(CompletionRequest.Model.GPT_3_5_TURBO_0125.getCode()).build();
//
//        ChatWebSocketRequest chatRequest = ChatWebSocketRequest.builder()
//                .chatContextList(messageContext)
//                .userId(111L)
//                .platform(AiPlatformEnum.OPEN_AI)
//                .sessionId("123456")
//                .chatModel(CompletionRequest.Model.GPT_3_5_TURBO_0125.getCode())
//                .build();
//        chatProcessorService.setChatRequest(chatRequest);
        return chatProcessorService.processSync();
    }
}
